"""
@author: 江同学呀
@file:  model_configuration.py
@date: 2025/2/23 17:17
@desc: 
    
"""
from typing import Dict, Literal, Optional

from espc.orm.model.base.base import _Base


class RegressionModelConfiguration(_Base):
    """
    分类模型的配置选项
    https://www.elastic.co/guide/en/elasticsearch/reference/7.17/search-aggregations-pipeline-inference-bucket-aggregation.html#inference-agg-regression-opt

    :param num_top_feature_importance_values:
        (Optional, integer) Specifies the maximum number of feature importance values per document. By default, it is
        zero and no feature importance calculation occurs.
        （可选，整数）指定每个文档的特征重要性值的最大数量。默认情况下，它为零，并且不会进行特征重要性计算。
    """

    def __init__(self, num_top_feature_importance_values: int = None):
        super().__init__()
        self._num_top_feature_importance_values: int = num_top_feature_importance_values
        return

    def _build(self) -> Dict:
        body: Dict = {}
        if self._num_top_feature_importance_values is not None:
            body["num_top_feature_importance_values"] = self._num_top_feature_importance_values
        return body


class ClassificationModelConfiguration(_Base):
    """
    分类模型的配置选项
    https://www.elastic.co/guide/en/elasticsearch/reference/7.17/search-aggregations-pipeline-inference-bucket-aggregation.html#inference-agg-classification-opt

    :param num_top_classes:
        (Optional, integer) Specifies the number of top class predictions to return. Defaults to 0.
        （可选，整数）指定要返回的顶级预测的数量。默认为 0。
    :param num_top_feature_importance_values:
        (Optional, integer) Specifies the maximum number of feature importance values per document. By default, it is
        zero and no feature importance calculation occurs.
        （可选，整数）指定每个文档的特征重要性值的最大数量。默认情况下，它为零，并且不会进行特征重要性计算。
    :param prediction_field_type:
        (Optional, string) Specifies the type of the predicted field to write. Acceptable values are: string, number,
        boolean. When boolean is provided 1.0 is transformed to true and 0.0 to false.
        （可选，字符串）指定要写入的预测字段的类型。可接受的值为 ： string， number， boolean.提供布尔值时，1.0 将转换为 true，0.0 将转
        换为 false。
    """

    def __init__(
            self, num_top_classes: int = None, num_top_feature_importance_values: int = None,
            prediction_field_type: Literal["string", "number", "boolean"] = None
    ):
        super().__init__()
        self._num_top_classes: int = num_top_classes
        self._num_top_feature_importance_values: int = num_top_feature_importance_values
        self._prediction_field_type: Optional[Literal["string", "number", "boolean"]] = prediction_field_type
        return

    def _build(self) -> Dict:
        body: Dict = {}
        if self._num_top_classes is not None:
            body["num_top_classes"] = self._num_top_classes
        if self._num_top_feature_importance_values is not None:
            body["num_top_feature_importance_values"] = self._num_top_feature_importance_values
        if self._prediction_field_type:
            body["prediction_field_type"] = self._prediction_field_type
        return body


